/* Copyright © 2024 Daniel Semyonov <daniel@dsemy.com>

   Permission is hereby granted, free of charge, to any person obtaining a
   copy of this software and associated documentation files (the "Software"),
   to deal in the Software without restriction, including without limitation
   the rights to use, copy, modify, merge, publish, distribute, sublicense,
   and/or sell copies of the Software, and to permit persons to whom the
   Software is furnished to do so, subject to the following conditions:

   The above copyright notice and this permission notice (including the next
   paragraph) shall be included in all copies or substantial portions of the
   Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
   DEALINGS IN THE SOFTWARE. */

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <wayland-client.h>
#include <wayland-client-protocol.h>

#include "wlr-output-power-management-unstable-v1-client-protocol.h"

#define WLR_DPMS_QUERY	4
#define WLR_DPMS_UNDEF	3
#define WLR_DPMS_ON	2
#define WLR_DPMS_OFF	1
#define WLR_DPMS_TOGGLE	0

struct power_state {
     uint32_t		 operation;
     struct wl_output	*output;
     char		*name;
};

struct manager_state {
     uint32_t					  operation;
     struct zwlr_output_power_manager_v1	 *manager;
     struct wl_output				**outputs;
     unsigned long				  outputs_n;
};

struct output_data {
     struct manager_state	 *state;
     int			  argc;
     char			**argv;
};

static void noop(void) {}

static void die(int status, char *message) {
     printf("ERROR: %s\n", message);
     exit(status);
}

static void power_free(void *data, struct zwlr_output_power_v1 *power) {
     struct power_state *state = (struct power_state *)data;
     zwlr_output_power_v1_destroy(power);
     wl_output_release(state->output);
     if (state->name) free(state->name);
     free(state);
}

static void power_mode(void *data, struct zwlr_output_power_v1 *power,
                 uint32_t mode) {
     struct power_state *state = (struct power_state *)data;
     if (state->operation == WLR_DPMS_QUERY) {
          printf("%s %s\n", state->name, mode ? "on" : "off");
     } else {
          zwlr_output_power_v1_set_mode(
              power,
              state->operation % WLR_DPMS_UNDEF ? --state->operation : !mode
          );
     }
     power_free(data, power);
}

static const struct zwlr_output_power_v1_listener plistener = {
     .mode = power_mode,
     .failed = power_free,
};

void output_register(struct manager_state *state,
                     struct wl_output *output, const char *name) {
     struct power_state *pstate = malloc(sizeof(struct power_state));
     if (!pstate) die(1, "output power state memory allocation failed");

     pstate->name = NULL;
     if (state->operation == WLR_DPMS_QUERY) {
          if (!name) die(4, "failed to get output name");
          pstate->name = malloc(sizeof(char) * (strlen(name) + 1));
          if (!pstate->name) die(1, "output name memory allocation failed");
          strcpy(pstate->name, name);
     }

     pstate->operation = state->operation;

     zwlr_output_power_v1_add_listener(
         zwlr_output_power_manager_v1_get_output_power(
             state->manager, (pstate->output = output)
         ),
         &plistener, pstate
     );
}

static void registry_global(void *data, struct wl_registry *registry,
                            uint32_t name, const char *interface, uint32_t _) {
     struct manager_state *state = (struct manager_state *)data;

     if (!strcmp(interface, zwlr_output_power_manager_v1_interface.name)) {
          state->manager = wl_registry_bind(
              registry, name,
              &zwlr_output_power_manager_v1_interface, 1
          );
     } else if (!strcmp(interface, wl_output_interface.name)) {
          state->outputs = realloc(state->outputs, ++state->outputs_n);
          if (!state->outputs) die(1, "outputs memory reallocation failed");
          state->outputs[state->outputs_n - 1] = wl_registry_bind(
                                                     registry, name,
                                                     &wl_output_interface, 4
                                                 );
     }
}

static const struct wl_registry_listener rlistener = {
     .global		= registry_global,
     .global_remove	= noop,
};

static void output_name (void *data, struct wl_output *output,
                         const char *name) {
        struct output_data *odata = (struct output_data *)data;

        if (odata->state->operation == WLR_DPMS_QUERY && odata->argc < 3) {
             output_register(odata->state, output, name);
             return;
        }

        for (int i = odata->state->operation == WLR_DPMS_UNDEF ? 1 : 2;
             i < odata->argc; i++)
             if (!strcmp(odata->argv[i], name)) {
                  output_register(odata->state, output, name);
                  return;
             }
        odata->state->outputs_n--;
}

static const struct wl_output_listener olistener = {
        .name		= output_name,
        .geometry	= noop,
        .mode		= noop,
        .scale		= noop,
        .description	= noop,
        .done		= noop,
};

static int names_required(int ac, uint32_t op) {
     return op == WLR_DPMS_QUERY || ac > 2 || (op == WLR_DPMS_UNDEF && ac > 1);
}

int main(int argc, char *argv[]) {
     struct manager_state *state = malloc(sizeof(struct manager_state));
     if (!state) die(1, "power manager state memory allocation failed");
     state->manager = NULL;
     state->outputs = NULL;
     state->outputs_n = 0;
     state->operation = WLR_DPMS_UNDEF;
     if (argc > 1) {
          if (!strcmp(argv[1], "on"))		state->operation = WLR_DPMS_ON;
          else if (!strcmp(argv[1], "off"))	state->operation = WLR_DPMS_OFF;
          else if (!strcmp(argv[1], "toggle"))	state->operation = WLR_DPMS_TOGGLE;
          else if (!strcmp(argv[1], "query"))	state->operation = WLR_DPMS_QUERY;
     }

     struct wl_display *display = wl_display_connect(NULL);
     if (!display) die(2, "display connection failed");
     struct wl_registry *registry = wl_display_get_registry(display);
     if (!registry) die(3, "display registry connection failed");
     wl_registry_add_listener(registry, &rlistener, state);
     wl_display_dispatch(display); /* bind manager and collect outputs */
     wl_registry_destroy(registry);

     struct output_data *data = NULL;
     if (names_required(argc, state->operation)) {
          data = malloc(sizeof(struct output_data));
          if (!data) die(1, "output data memory allocation failed");
          data->state = state;
          data->argc = argc;
          data->argv = argv;
     }

     for (unsigned long i = 0; i < state->outputs_n; i++)
          if (names_required(argc, state->operation))
               wl_output_add_listener(state->outputs[i], &olistener, data);
          else output_register(state, state->outputs[i], NULL);
     if (names_required(argc, state->operation))
          wl_display_dispatch(display); /* collect names */
     zwlr_output_power_manager_v1_destroy(state->manager);

     if (state->outputs_n) {
          wl_display_dispatch(display);  /* register power listeners */
          wl_display_roundtrip(display); /* set power modes */
     } else {
          fprintf(stderr, "WARNING: no specified output found\n");
     }

     wl_display_disconnect(display);
     free(state->outputs);
     free(state);
}
